structure Prfgen =
struct
  structure P = Proof
  structure Asm = Assem
  structure A = Assert

  fun mapToAddr (id) = 0 (* cheat compiler! *)
  fun reg2exp(reg) = 
      let fun handlereg(Asm.Eax) = A.Eax
	    | handlereg(Asm.Ebx) = A.Ebx
	    | handlereg(Asm.Ecx) = A.Ecx
	    | handlereg(Asm.Edx) = A.Edx
	    | handlereg(Asm.Esi) = A.Esi
	    | handlereg(Asm.Edi) = A.Edi
	    | handlereg(Asm.Esp) = A.Esp
	    | handlereg(Asm.Ebp) = A.Ebp
      in handlereg(reg)
      end

(* update State: state*instr-> state *)
  fun upState instr st= 
      let fun rf(state:A.state) = #1 state
	  fun hp(state:A.state) = #2 state
	  fun mchinstr(Asm.Addr(r1,r2)) = 
	      let val updexp2 = A.Plus(A.Sel(rf(st),reg2exp(r2)),
				       A.Sel(rf(st),reg2exp(r1)))
		  val st' = (A.Upd(rf(st),reg2exp(r2),updexp2),hp(st))
	      in st'
	      end
	    | mchinstr(Asm.Addz(imm,r2)) = 
	      let val updexp2 = A.Plus(A.Sel(rf(st),reg2exp(r2)),
				       A.IntNum(imm))
		  val st' = (A.Upd(rf(st),reg2exp(r2),updexp2),hp(st))
	      in st'
	      end
	    | mchinstr(Asm.Subr(r1,r2)) = 
	      let val updexp2 = A.Minus(A.Sel(rf(st),reg2exp(r2)),
				       A.Sel(rf(st),reg2exp(r1)))
		  val st' = (A.Upd(rf(st),reg2exp(r2),updexp2),hp(st))
	      in st'
	      end
	    | mchinstr(Asm.Subz(imm,r2)) = 
	      let val updexp2 = A.Minus(A.Sel(rf(st),reg2exp(r2)),
				       A.IntNum(imm))
		  val st' = (A.Upd(rf(st),reg2exp(r2),updexp2),hp(st))
	      in st'
	      end
	    | mchinstr(Asm.Mulr(r1,r2)) = 
	      let val updexp2 = A.Times(A.Sel(rf(st),reg2exp(r2)),
				       A.Sel(rf(st),reg2exp(r1)))
		  val st' = (A.Upd(rf(st),reg2exp(r2),updexp2),hp(st))
	      in st'
	      end
	    | mchinstr(Asm.Mulz(imm,r2)) = 
	      let val updexp2 = A.Times(A.Sel(rf(st),reg2exp(r2)),
				       A.IntNum(imm))
		  val st' = (A.Upd(rf(st),reg2exp(r2),updexp2),hp(st))
	      in st'
	      end
	    | mchinstr(Asm.Push(r1)) = 
	      let val rfupdexp2 = A.Minus(A.Sel(rf(st),A.Esp),A.IntNum(4))
		  val rf' = A.Upd(rf(st),A.Esp,rfupdexp2)
		  val hp' = A.Upd(hp(st),A.Sel(rf(st),A.Esp),
				  A.Sel(rf(st),reg2exp(r1)))
		  val st' = (rf',hp')
	      in st'
	      end
	    | mchinstr(Asm.Pop(r1)) = 
	      let val updesp = A.Upd(rf(st), A.Esp,
				     A.Plus(A.Sel(rf(st),A.Esp),A.IntNum(4)))
		  val rf' = A.Upd(updesp,reg2exp(r1),
				  A.Sel(hp(st),A.Sel(rf(st),A.Esp)))
		  val st' = (rf',hp(st))
	      in st'
	      end
	    | mchinstr(Asm.Xchg(r1,r2)) = 
	      let val rf1 = A.Upd(rf(st),reg2exp(r1),A.Sel(rf(st),reg2exp(r2)))
		  val st' = (A.Upd(rf1,reg2exp(r2),A.Sel(rf(st),reg2exp(r1))),hp(st))
	      in st'
	      end
	    | mchinstr(Asm.Leal(str,r1)) = 
	      let val st' = (A.Upd(rf(st),reg2exp(r1),A.Addr(mapToAddr(str))),hp(st))
	      in st'
	      end
	    | mchinstr(Asm.Leab(ofaddr,r1,r2)) = (*WIP*)
	      let val addr = A.Plus(A.Sel(rf(st),reg2exp(r1)),
				    A.IntNum(ofaddr))
		  val st' = (A.Upd(rf(st),reg2exp(r2),addr),hp(st))
	      in st'
	      end
	    | mchinstr(Asm.Movr(r1,r2)) = 
	      let val st' = (A.Upd(rf(st),reg2exp(r2),A.Sel(rf(st),reg2exp(r1)))
			     ,hp(st))
	      in st'
	      end
	    | mchinstr(Asm.Movz(imm,r2)) = 
	      let val st' = (A.Upd(rf(st),reg2exp(r2),A.IntNum(imm)),hp(st))
	      in st'
	      end
	    | mchinstr(Asm.Movld(imm,r1,r2)) = 
	      let val updexp1 = A.Sel(hp(st),A.Plus(A.Sel(rf(st),reg2exp(r1)),
						    A.IntNum(imm)))
		  val st' = (A.Upd(rf(st),reg2exp(r2),updexp1),hp(st))
	      in st'
	      end
	    | mchinstr(Asm.Movst(r1,imm,r2)) = 
	      let val updexp1 = A.Plus(A.Sel(rf(st),reg2exp(r2)),A.IntNum(imm))
		  val hp' = A.Upd(hp(st),updexp1,A.Sel(rf(st),reg2exp(r1)))
		  val st' = (rf(st),hp')
	      in st'
	      end
            | mchinstr(Asm.Ret) =
              let val rf' = A.Upd(rf(st), A.Esp,
				 A.Plus(A.Sel(rf(st),A.Esp),A.IntNum(4)))
                  val st' = (rf', hp(st))
              in  st'
              end
	    | mchinstr(_) = st
      in mchinstr(instr)
      end


  fun prf_gen pre (cd_b, prf_b) post cb specs vcname = 
      let fun gen ((Asm.Jmp(l))::cb') =
              let val postopt = A.getspec l specs
                  val post = case (!postopt)
                              of  NONE => (raise Fail ("Impossible! No such label: "^l); (fn st:A.state => A.True))
                                | SOME(a) => a
                  val st = upState (Asm.Jmp(l))
                  val a = (fn st_init:A.state => post (st st_init)) 
              in  case cb'
                   of  [] => ([Asm.Jmp(l)],P.WF_jmp (l,(pre, post, vcname)))
                     | _ => let val new_prf_b = P.WF_jmp (l,(a, post, ""))
                                val new_cd_b = [Asm.Jmp(l)]
                            in  prf_gen pre (new_cd_b, new_prf_b) a cb' specs vcname
                            end 
              end
            | gen ((Asm.Addr(r1, r2))::cb') =
              let val st = upState (Asm.Addr(r1,r2))
                  val a = (fn st_init:A.state => post (st st_init))
              in  case cb'
                   of  [] => ((Asm.Addr(r1,r2))::cd_b,(P.WF_addr (r1, r2, (cd_b, prf_b), (pre, post, vcname))))
                     | _ => let val new_prf_b = P.WF_addr (r1, r2, (cd_b, prf_b), (a, post, "")) 
                                val new_cd_b = (Asm.Addr(r1,r2))::cd_b
                            in  prf_gen pre (new_cd_b, new_prf_b) a cb' specs vcname
                            end
              end
            | gen (i::cb') = (raise Fail "work in progress!"; (i::cd_b, P.WF_trivial))
            | gen [] = (raise Fail "impossible!"; (cd_b, P.WF_trivial))
      in  gen cb
      end


  fun prf_gen_no_pre (cd_b, prf_b) post cb specs lab = 
      let fun gen ((Asm.Jmp(l))::cb') =
              let val postopt = A.getspec l specs
                  val post = case !postopt
                              of  NONE => (raise Fail ("Impossible! No such label: "^l); (fn st:A.state => A.True))
                                | SOME(a) => a
                  val st = upState (Asm.Jmp(l))
                  val a = (fn st_init:A.state => post (st st_init)) 
              in  case cb'
                   of  [] => let val pre = A.getspec lab specs
                                 val _ = pre := (SOME a) 
                             in ([Asm.Jmp l],P.WF_jmp (l,(a, post, "")))
                             end
                     | _ => let val new_prf_b = P.WF_jmp (l, (a, post, ""))
                                val new_cd_b = [Asm.Jmp l]
                            in  prf_gen_no_pre (new_cd_b, new_prf_b) a cb' specs lab
                            end 
              end
            | gen ((Asm.Addr(r1, r2))::cb') =
              let val st = upState (Asm.Addr(r1,r2))
                  val a = (fn st_init:A.state => post (st st_init))
              in  case cb'
                   of  [] =>  let val pre = A.getspec lab specs
                                  val _ = pre := (SOME a) 
                             in (((Asm.Addr (r1, r2))::cd_b), P.WF_addr (r1, r2, (cd_b, prf_b), (a, post, "")))
                             end
                     | _ => let val new_prf_b = P.WF_addr (r1, r2, (cd_b, prf_b), (a, post, ""))
                                val new_cd_b = (Asm.Addr (r1, r2))::cd_b 
                            in  prf_gen_no_pre (new_cd_b, new_prf_b) a cb' specs lab
                            end
              end
            | gen (i::cb') = (print "work in progress!\n"; (i::cd_b, P.WF_trivial)) (*WIP*)
            | gen [] = (raise Fail "impossible!"; (cd_b, P.WF_trivial))
      in  gen cb
      end

  fun get_cb code lab =
      let fun getlab [] = (raise Fail "no such lab in codes!"; [])
            | getlab ((Asm.Label(l, a))::code') = if (l = lab)
                                                  then code'
                                                  else getlab code'
            | getlab (_::code') = getlab code'
          val begin = getlab code
          val cb = ref []
          fun get [] = ()
            | get (Asm.Label(_,_)::code') = ()
            | get (i::code') = (cb := i::(!cb); (get code'))
      in  get begin;
          !(cb)
      end

  fun prf_list_gen_1 spec_index code specs =
      let val prf_list = ref []
          fun gen (l::speci) = 
              let val cb = get_cb code l
                  val l1 = case cb 
                            of (Asm.Jmp(l2)::cb') => l2
                             | (Asm.Ret::cb')=> "post_condition" (*WIP*)
                             | _ => (raise Fail "impossible, last instruction is not a jmp!"; "bogus")
                  fun eql x = if (x = l1)
                              then true
                              else false
              in  if (List.exists eql (l::speci))
                  then gen (speci@[l])
                  else let val post' = A.getspec l1 specs
                           val post = case (!post') 
                                       of  NONE => (raise Fail ("impossible! no such specification: "^l1); fn s:A.state => A.True)
                                         | SOME (a) => a
                           val proof = prf_gen_no_pre ([], P.WF_trivial) post cb specs l
                       in  prf_list := ((l, proof):: !prf_list);
                           gen speci
                       end
              end
            | gen [] = ()
      in  gen spec_index;
          !prf_list
      end

  fun prf_list_gen_2 spec_index code specs =
      let val prf_list = ref []
          fun gen (l::speci) = 
              let val cb = get_cb code l
                  val l1 = case cb 
                            of (Asm.Jmp(l2)::cb') => l2
                             | (Asm.Ret::cb')=> "post_condition" (*WIP*)
                             | _ => (raise Fail "impossible, last instruction is not a jmp!"; "bogus")
              in  let val post' = A.getspec l1 specs
                      val post = case (!post') 
                                  of  NONE => (raise Fail ("impossible! no such specification: "^l1); fn s:A.state => A.True)
                                    | SOME (a) => a
                      val pre' = A.getspec l specs
                      val pre = case (!pre')
                                 of  NONE => (raise Fail ("impossible! no such specification: "^l); fn s:A.state => A.True)
                                   | SOME (a) => a
                      val proof = prf_gen pre ([], P.WF_trivial) post cb specs l
                  in  prf_list := ((l, proof):: !prf_list);
                      gen speci
                  end
              end
            | gen [] = ()
      in  gen spec_index;
          !prf_list
      end

end
